"""
网络连接管理服务 - 支持TCP和UDP
"""
import socket
import select
import time
from typing import Dict, Optional
from models import NetConfig, ConnectionStats, APIResponse

class NetConnectionManager:
    """网络连接管理器 - 支持TCP和UDP"""
    
    def __init__(self):
        self.connections: Dict[str, dict] = {}
    
    def create_connection(self, config: NetConfig) -> APIResponse:
        """
        创建网络连接
        
        Args:
            config: 网络配置（支持TCP和UDP）
            
        Returns:
            APIResponse: 连接结果
        """
        # 验证配置
        if config.protocol not in ["tcp", "udp"]:
            return APIResponse(
                success=False,
                message=f"不支持的协议类型: {config.protocol}"
            )
        
        if config.port < 1 or config.port > 65535:
            return APIResponse(
                success=False,
                message=f"端口号必须在1-65535范围内，当前值: {config.port}"
            )
        
        connection_id = f"{config.protocol}:{config.host}:{config.port}"
        
        try:
            if config.protocol == "tcp":
                sock = self._create_tcp_connection(config)
            else:  # udp
                sock = self._create_udp_connection(config)
                
            self.connections[connection_id] = {
                "socket": sock,
                "config": config.dict(),
                "connected": True,
                "stats": ConnectionStats().dict(),
                "clients": [] if config.mode == "server" and config.protocol == "tcp" else None,
                "remote_address": None if config.mode == "server" else (config.host, config.port)
            }
            
            message = f"成功创建{config.protocol.upper()}连接: {config.host}:{config.port}" if config.mode == "client" else f"{config.protocol.upper()}服务器已启动，监听 {config.host}:{config.port}"
            
            response_data = {
                "success": True,
                "message": message,
                "data": {"connection_id": connection_id},
                "connection_id": connection_id  # 添加顶级字段兼容性
            }
            
            return APIResponse(**response_data)
            
        except Exception as e:
            return APIResponse(
                success=False,
                message=f"创建连接失败: {str(e)}"
            )
            
    def _create_tcp_connection(self, config: NetConfig) -> socket.socket:
        """创建TCP连接"""
        sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
        sock.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
        
        if config.mode == "client":
            sock.settimeout(3.0)  # 客户端连接超时设置为3秒
            sock.connect((config.host, config.port))
            sock.settimeout(0.1)  # 连接成功后设置为非阻塞模式
        else:  # server mode
            sock.bind((config.host, config.port))
            sock.listen(5)
            sock.settimeout(0.1)  # 设置非阻塞模式
        
        return sock
        
    def _create_udp_connection(self, config: NetConfig) -> socket.socket:
        """创建UDP连接"""
        sock = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
        sock.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
        
        if config.mode == "server":
            sock.bind((config.host, config.port))
        
        sock.settimeout(0.1)  # 设置短超时以避免阻塞
        return sock
        
    def disconnect(self, connection_id: str) -> APIResponse:
        """
        断开网络连接
        
        Args:
            connection_id: 连接ID
            
        Returns:
            APIResponse: 断开结果
        """
        if connection_id not in self.connections:
            return APIResponse(
                success=False,
                message="连接不存在"
            )
        
        try:
            conn = self.connections[connection_id]
            
            # 关闭客户端连接（如果是TCP服务器模式）
            if conn.get("clients"):
                for client in conn["clients"]:
                    try:
                        client.close()
                    except:
                        pass
            
            # 关闭主socket
            conn["socket"].close()
            del self.connections[connection_id]
            
            return APIResponse(
                success=True,
                message="连接已断开"
            )
            
        except Exception as e:
            return APIResponse(
                success=False,
                message=f"断开连接失败: {str(e)}"
            )
            
    def send_data(self, connection_id: str, data: bytes) -> APIResponse:
        """
        发送网络数据
        
        Args:
            connection_id: 连接ID
            data: 要发送的数据
            
        Returns:
            APIResponse: 发送结果
        """
        if connection_id not in self.connections:
            return APIResponse(
                success=False,
                message="连接不存在"
            )
        
        try:
            conn = self.connections[connection_id]
            sock = conn["socket"]
            
            if conn["config"]["protocol"] == "tcp":
                if conn["config"]["mode"] == "client":
                    # 确保socket有效且连接状态正常
                    try:
                        sock.send(data)
                    except BrokenPipeError:
                        conn["connected"] = False
                        return APIResponse(
                            success=False,
                            message="连接已断开（BrokenPipeError）"
                        )
                    except ConnectionResetError:
                        conn["connected"] = False
                        return APIResponse(
                            success=False,
                            message="连接被重置（ConnectionResetError）"
                        )
                else:
                    # TCP服务器模式
                    clients = conn.get("clients", [])
                    if not clients:
                        return APIResponse(
                            success=False,
                            message="没有连接的客户端"
                        )
                    
                    # 发送给所有客户端，并处理可能的连接错误
                    success_clients = 0
                    for client in clients:
                        try:
                            client.send(data)
                            success_clients += 1
                        except Exception:
                            # 客户端连接错误，标记为断开
                            try:
                                client.close()
                            except:
                                pass
                            # 注意：不在这里移除客户端，避免修改正在迭代的列表
                    
                    # 移除已断开的客户端（在单独的循环中）
                    conn["clients"] = [client for client in clients if client.fileno() != -1]
                    
                    if success_clients == 0:
                        return APIResponse(
                            success=False,
                            message="所有客户端连接已断开"
                        )
            else:
                # UDP模式
                if conn.get("remote_address"):
                    try:
                        sock.sendto(data, conn["remote_address"])
                    except Exception as e:
                        return APIResponse(
                            success=False,
                            message=f"UDP发送失败: {str(e)}"
                        )
                else:
                    return APIResponse(
                        success=False,
                        message="UDP服务器模式下请先设置远程地址"
                    )
            
            # 更新统计信息
            conn["stats"]["send_count"] += 1
            conn["stats"]["send_bytes"] += len(data)
            
            return APIResponse(
                success=True,
                message="数据发送成功"
            )
            
        except Exception as e:
            # 捕获并记录所有异常
            print(f"发送数据异常: {str(e)}")
            # 更新连接状态
            if connection_id in self.connections:
                self.connections[connection_id]["stats"]["error_count"] += 1
            return APIResponse(
                success=False,
                message=f"发送失败: {str(e)}"
            )
            
    def read_data(self, connection_id: str) -> Optional[bytes]:
        """
        读取网络数据
        
        Args:
            connection_id: 连接ID
            
        Returns:
            Optional[bytes]: 读取的数据，如果没有数据则返回None
        """
        if connection_id not in self.connections:
            return None
        
        try:
            conn = self.connections[connection_id]
            sock = conn["socket"]
            
            # 使用select进行IO多路复用，避免阻塞
            if conn["config"]["protocol"] == "tcp":
                if conn["config"]["mode"] == "client":
                    # TCP客户端模式：检查是否有数据可读
                    try:
                        # 使用select检查是否有数据可读
                        readable, _, exceptional = select.select([sock], [], [sock], 0.01)  # 10ms超时
                        if sock in exceptional:
                            conn["connected"] = False
                            return None
                        
                        if sock in readable:
                            data = sock.recv(4096)
                            if data:
                                # 更新统计信息
                                conn["stats"]["receive_bytes"] += len(data)
                                conn["stats"]["receive_count"] += 1
                                return data
                            else:
                                # 连接关闭
                                conn["connected"] = False
                                return None
                    except socket.timeout:
                        pass
                    except Exception as e:
                        print(f"TCP客户端读取数据异常: {str(e)}")
                        conn["stats"]["error_count"] += 1
                        conn["connected"] = False
                        return None
                else:
                    # TCP服务器模式：使用select处理多客户端
                    try:
                        # 构建要监听的套接字列表
                        sockets_to_watch = [sock]  # 服务器套接字
                        if conn.get("clients"):
                            sockets_to_watch.extend(conn["clients"])
                        
                        # 使用select检查是否有活动
                        readable, _, exceptional = select.select(sockets_to_watch, [], sockets_to_watch, 0.01)  # 10ms超时
                        
                        # 处理异常套接字
                        for s in exceptional:
                            if s == sock:
                                # 服务器套接字异常
                                conn["connected"] = False
                                return None
                            elif conn.get("clients") and s in conn["clients"]:
                                # 客户端套接字异常，移除它
                                try:
                                    s.close()
                                except:
                                    pass
                                if s in conn["clients"]:
                                    conn["clients"].remove(s)
                        
                        # 处理可读套接字
                        for s in readable:
                            if s == sock:
                                # 有新连接
                                try:
                                    client_sock, addr = s.accept()
                                    client_sock.settimeout(0.1)
                                    conn["clients"].append(client_sock)
                                    print(f"新客户端连接: {addr}")
                                except Exception:
                                    pass
                            elif conn.get("clients") and s in conn["clients"]:
                                # 从客户端接收数据
                                try:
                                    data = s.recv(4096)
                                    if data:
                                        conn["stats"]["receive_bytes"] += len(data)
                                        conn["stats"]["receive_count"] += 1
                                        return data
                                    else:
                                        # 客户端关闭连接
                                        try:
                                            s.close()
                                        except:
                                            pass
                                        if s in conn["clients"]:
                                            conn["clients"].remove(s)
                                except Exception:
                                    # 客户端连接错误，移除它
                                    try:
                                        s.close()
                                    except:
                                        pass
                                    if s in conn["clients"]:
                                        conn["clients"].remove(s)
                    except Exception as e:
                        print(f"TCP服务器读取数据异常: {str(e)}")
                        conn["stats"]["error_count"] += 1
                        return None
            else:
                # UDP模式
                try:
                    # 使用select检查是否有数据可读
                    readable, _, exceptional = select.select([sock], [], [sock], 0.01)  # 10ms超时
                    if sock in exceptional:
                        conn["connected"] = False
                        return None
                    
                    if sock in readable:
                        data, addr = sock.recvfrom(4096)
                        if data:
                            # 更新统计信息
                            conn["stats"]["receive_bytes"] += len(data)
                            conn["stats"]["receive_count"] += 1
                            
                            # 对于UDP服务器模式，记录最后一个发送数据的客户端地址
                            if conn["config"]["mode"] == "server":
                                conn["remote_address"] = addr
                            
                            return data
                except socket.timeout:
                    pass
                except Exception as e:
                    print(f"UDP读取数据异常: {str(e)}")
                    conn["stats"]["error_count"] += 1
                    return None
            
        except Exception as e:
            # 更新错误统计
            if connection_id in self.connections:
                self.connections[connection_id]["stats"]["error_count"] += 1
                print(f"读取数据异常: {str(e)}")
        
        return None

    def get_connection_status(self, connection_id: str) -> Optional[Dict]:
        """
        获取连接状态
        
        Args:
            connection_id: 连接ID
            
        Returns:
            Optional[Dict]: 连接状态信息
        """
        if connection_id in self.connections:
            conn = self.connections[connection_id]
            status = {
                "connected": conn["connected"],
                "protocol": conn["config"].get("protocol", "tcp"),
                "mode": conn["config"].get("mode", "client"),
                "host": conn["config"].get("host", ""),
                "port": conn["config"].get("port", 0),
                "stats": conn["stats"]
            }
            
            # 添加服务器模式的客户端信息
            if conn.get("clients"):
                status["clients"] = len(conn["clients"])
            
            # 添加UDP模式的远程地址
            if conn.get("remote_address"):
                status["remote_address"] = conn["remote_address"]
            
            return status
        return None
    
    def get_connection(self, connection_id: str) -> Optional[Dict]:
        """
        获取连接信息
        
        Args:
            connection_id: 连接ID
            
        Returns:
            Optional[Dict]: 连接信息字典
        """
        return self.connections.get(connection_id)

# 实例化网络连接管理器
net_manager = NetConnectionManager()